import numpy as np
import torch

if __name__ == "__main__":

    # y = np.array([0, 1, 1, 0, 0, 0])
    # y_hat = np.array([0.9, 0.8, 0.7, 0.6, 0.5, 0.4])
    # print(recall_at_k(y, y_hat, 3), "true is: ", 1)
    # print(precision_at_k(y, y_hat, 3), "true is: ", 2/3)
    # print(ap_at_k(y, y_hat, 3), "true is: ", 7/12)
    # print(auc_at_k(y, y_hat, 4), "true is: ", 0.5)
    # print(ndcg_at_k(y, y_hat, 3), "true is: ", 0.6934264036172708)
    # print(mrr_at_k(y, y_hat, 3), "true is: ", 1/2)
    # print(metrics_at_Ks(y, y_hat, [1, 2, 3, 4, 10, 15]))

    # batch
    y = np.array([[0, 1, 1, 0, 1, 0],
                  [0, 1, 1, 0, 1, 0],
                  [1, 0, 1, 1, 0, 1]])
    y_hat = np.array([[0.9, 0.8, 0.7, -0.6, 0, -0.4],
                      [0.9, 0.8, 0.7, -0.6, 0.5, -0.4],
                      [0.9, 0.8, 0.7, -0.6, 0.5, -0.4]])
    # print(torch.mul(torch.Tensor(y), torch.arange(6)))
    # print(torch.dot(torch.Tensor(y), torch.arange(6)))
    # print(torch.heaviside(torch.Tensor(y_hat), torch.tensor([0.5])))
    # print(torch.sigmoid(torch.Tensor(y_hat)))
    # print(torch.pow(2, torch.arange(10)))
    a = set([1, 2])
    b = set([3, 4])
    # print(a + b)
    print(a & b)
    c = a | b
    print(a & c)